{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4cba3ea5",
   "metadata": {},
   "outputs": [],
   "source": [
    "remote = \"https://raw.githubusercontent.com/nansencenter/DA-tutorials\"\n",
    "!wget -qO- {remote}/master/notebooks/resources/colab_bootstrap.sh | bash -s"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "08888704",
   "metadata": {},
   "outputs": [],
   "source": [
    "from resources import show_answer, interact,  import_from_nb, nonchalance, colorbar\n",
    "%matplotlib inline\n",
    "import numpy as np\n",
    "import matplotlib as mpl\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy.random as rnd\n",
    "import scipy.linalg as sla\n",
    "from mpl_tools.place import freshfig\n",
    "plt.ion();"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a42a0085",
   "metadata": {},
   "source": [
    "# T7 - Spatial statistics (\"geostatistics\") & kriging\n",
    "\n",
    "In T4 and T5, we performed state estimation for scalar and vector (length-4) noisy, linear AR(1) processes,\n",
    "while T6 illustrated several nonlinear and chaotic dynamics, including the 1D (length-40) Lorenz-96 system.\n",
    "We now turn to estimating static (non-dynamic) 1D, 2D, and 3D spatial fields\n",
    "using the method known (in geostatistics) as **kriging**,\n",
    "**radial basis function (RBF)** interpolation (in mathematics),\n",
    "and **Gaussian process (GP)** regression (in machine learning).\n",
    "$\n",
    "\\newcommand{\\mat}[1]{{\\mathbf{{#1}}}}\n",
    "\\newcommand{\\vect}[1]{{\\mathbf{#1}}}\n",
    "\\newcommand{\\Expect}[0]{\\mathbb{E}}\n",
    "\\newcommand{\\NormDist}{\\mathscr{N}}\n",
    "\\newcommand{\\trsign}{{\\mathsf{T}}}\n",
    "\\newcommand{\\tr}{^{\\trsign}}\n",
    "$\n",
    "\n",
    "Once again, randomness will serve as a stand-in for uncertainty and will generate our synthetic experiments.\n",
    "Conceptually, little is new – once discretized and flattened, spatial fields are simply long vectors of unknowns, which we can estimate using covariances very similarly as in the Kalman filter.\n",
    "The challenge lies in their scale: such long vectors lead to large and unwieldy covariance matrices.\n",
    "Moreover, rather than treating the covariance matrix as merely a symmetric table of numbers,\n",
    "we will explicitly model it through spatial relationships and distances\n",
    "between field elements. Such covariance modeling will also be useful for the EnKF in that they can encode initial priors, and the experience gained will help in understanding covariance localization.\n",
    "\n",
    "As always, we discretize the domains.\n",
    "The following sets up a grid/lattice for fields on a 1D domain."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a5292eae",
   "metadata": {},
   "outputs": [],
   "source": [
    "grid1D = np.linspace(0, 1, 21)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "280f6d4a",
   "metadata": {},
   "source": [
    "The following sets up a 2D grid."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "baf5431e",
   "metadata": {},
   "outputs": [],
   "source": [
    "grid2x, grid2y = np.meshgrid(grid1D, grid1D)\n",
    "grid2x.shape  # `grid2y` has same shape"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "93497302",
   "metadata": {},
   "source": [
    "Now we \"flatten\" (a.k.a. \"(un)ravel\", \"vectorize\", or \"string out\") this explicitly 2D grid of nodes into a list of 2D points."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "89116e62",
   "metadata": {},
   "outputs": [],
   "source": [
    "grid2D = np.dstack([grid2x, grid2y]).reshape((-1, 2))\n",
    "grid2D.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "438460b0",
   "metadata": {},
   "source": [
    "## Random fields and variograms\n",
    "\n",
    "Denote a *physical* location (i.e., a 1D/2D/3D coordinate) by $s$.\n",
    "A ***random field***, $X(s)$, is a function that takes random values at each point $s$.\n",
    "The near-synonymous term *regionalized variable* emphasizes the continuity of $s$.\n",
    "\n",
    "Random fields are commonly assumed to be ***stationary***,\n",
    "meaning that the dependence between $X(s_{1})$ and $X(s_{2})$\n",
    "for any two locations $s_1, s_2$ can be described solely in terms of their separation or,\n",
    "assuming isotropy (rotational symmetry), their distance $d = \\| s_{1} - s_{2} \\|$.\n",
    "\n",
    "For now, we assume that the mean $\\Expect X(s)$ is known and, by stationarity, constant,\n",
    "leaving the covariance as the most important descriptor.\n",
    "Again, by stationarity, this depends only on distance,\n",
    "so we can describe the full covariance of the field solely in terms of the\n",
    "\"(auto-)***covariance function***\": $$C(d) = \\mathbb{Cov}[X(s_{1}), X(s_{2})] \\,.$$\n",
    "<details style=\"border: 1px solid #aaaaaa; border-radius: 4px; padding: 0.5em 0.5em 0;\">\n",
    "  <summary style=\"font-weight: normal; font-style: italic; margin: -0.5em -0.5em 0; padding: 0.5em;\">\n",
    "    Geostatistics usually works with a reformulation of $C(d)$ called \"(semi-)<strong>variogram\"</strong>,\n",
    "    which flips the graph of $C(d)$ on its head ... (optional reading 🔍):\n",
    "    $$\\gamma(d) = C(0) - C(d)$$\n",
    "  </summary>\n",
    "\n",
    "  Not all functions are valid variograms; they must evidently be non-negative,\n",
    "  have the value $0$ at $d=0$ (why?),\n",
    "  and result in positive-definite covariance matrices.\n",
    "  A sufficient condition is that the corresponding covariance is convex and decays\n",
    "  (but does not reach) toward zero at infinity.\n",
    "  However, covariance functions that oscillate around zero (Matérn)\n",
    "  or have finite support (triangular) also exist.\n",
    "\n",
    "  A wider definition of the variogram also exists:\n",
    "  half (whence the optional \"semi\" prefix) the variance of the increment, i.e. of $X(s_{1}) - X(s_{2})$.\n",
    "  This definition only requires second-order stationarity of the *increments*,\n",
    "  which is a weaker requirement called \"intrinsic stationarity\",\n",
    "  and does not require or imply the existence of a covariance function $C(d)$.\n",
    "  Thus, it can be used to estimate infinite-variance (heavy-tail or Brownian-motion-like) processes,\n",
    "  but only if the method constrains the sum of its weights to one,\n",
    "  making the linear combination of the estimate \"authorized/allowable.\"\n",
    "  This condition coincides with that imposed by an unknown mean\n",
    "  and does not represent a severe restriction.\n",
    "  - - -\n",
    "</details>\n",
    "\n",
    "Turning things around, we can start by prescribing a *model* variogram (i.e., covariance)\n",
    "that we believe the random field has.\n",
    "Some of the most commonly used examples are implemented below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3ccf721a",
   "metadata": {},
   "outputs": [],
   "source": [
    "def variograms(power=1.5, transf01=\"expo\", Range=0.3, nugget=0, sill=1):\n",
    "    \"\"\"Create variogram (function) for the given parameters.\"\"\"\n",
    "    def vg(dists):\n",
    "        dists = np.asarray(dists)\n",
    "        dists = dists / Range\n",
    "        gamma = transforms01[transf01](dists**power)\n",
    "        gamma *= sill\n",
    "        gamma = np.where(dists != 0, nugget + (1 - nugget) * gamma, gamma)\n",
    "        return gamma\n",
    "    return vg"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d3ab629b",
   "metadata": {},
   "source": [
    "Here, `transforms01` monotonically decreasingly maps $[0, +\\infty]$ to $[0, 1]$ ."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "82064f96",
   "metadata": {},
   "outputs": [],
   "source": [
    "transforms01 = {\n",
    "    \"none\":  lambda x: x,               # identity\n",
    "    \"expo\":  lambda x: 1 - np.exp(-x),  # similar to Gaussian densities\n",
    "    \"hyper\": lambda x: 1 - 1 / (1 + x), # similar to Cauchy densities\n",
    "    \"clip\":  lambda x: x.clip(max=1)    # similar to \"hat\" function\n",
    "}\n",
    "vg_params = dict(transf01=list(transforms01), power=(0.01, 4, .01),\n",
    "                 Range=(.01, 4), nugget=(0, 1, .01), sill=(0.1, 5))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6c53d9c1",
   "metadata": {},
   "source": [
    "Below is a visual illustration."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0b49c807",
   "metadata": {},
   "outputs": [],
   "source": [
    "@interact(**vg_params)\n",
    "def plot_variograms(power=1.5, Range=0.3, nugget=0, sill=1):\n",
    "    fig, ax = plt.subplots(figsize=(6, 3))\n",
    "    ax.set_ylim(0, 1.05 * sill)\n",
    "    ax.set_xlim(left=0)\n",
    "    for i, transf01 in enumerate(transforms01):\n",
    "        vg = variograms(power, transf01, Range, nugget, sill)\n",
    "        ax.plot(grid1D, vg(grid1D), lw=5-i, color=f\"C{i+1}\", label=transf01)\n",
    "    plt.xlabel(\"distance/lag\")\n",
    "    ax.legend(title=\"transf 0→1\", loc=\"lower right\")\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "022629eb",
   "metadata": {},
   "source": [
    "Now that we're in possession of a valid variogram,\n",
    "we can compute a covariance matrix between any sets of points as follows."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "17026cbd",
   "metadata": {},
   "outputs": [],
   "source": [
    "def covar_modelled(coords, variogram):\n",
    "    C0 = variogram(np.inf)\n",
    "    assert C0 < np.inf, \"Not a valid covariance\"\n",
    "    dists = dist_euclid(coords, coords)\n",
    "    return C0 - variogram(dists)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f2ad003b",
   "metadata": {},
   "source": [
    "We also need to be able to compute pair-wise distances for each (discretized) location of our field.\n",
    "The following computes distances between two sets (not necessarily of same length) of points."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4bbfe6aa",
   "metadata": {},
   "outputs": [],
   "source": [
    "def dist_euclid(A, B):\n",
    "    \"\"\"l2-norm between each point (row) of A and B\"\"\"\n",
    "    # Like scipy.spatial.distance.pdist(A) + `squareform` if A==B\n",
    "    if A.ndim == 1:\n",
    "        return abs(A[:, None] - B)\n",
    "    d = A[:, None, :] - B\n",
    "    d2 = np.sum(d**2, axis=-1)\n",
    "    return np.sqrt(d2)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c3397d6d",
   "metadata": {},
   "source": [
    "## Gaussian fields\n",
    "\n",
    "<details style=\"border: 1px solid #aaaaaa; border-radius: 4px; padding: 0.5em 0.5em 0;\">\n",
    "  <summary style=\"font-weight: normal; font-style: italic; margin: -0.5em -0.5em 0; padding: 0.5em;\">\n",
    "    Gaussianity becomes even more useful when working with spatial fields ... (optional reading 🔍)\n",
    "  </summary>\n",
    "  Non-Gaussian random fields, such as those with Student-t or Laplace marginals, exist,\n",
    "  but Gaussianity makes sampling and conditioning uniquely\n",
    "  tractable, even for large vectors, since everything reduces to linear algebra.\n",
    "\n",
    "  Besides, although **Markov** random fields of any dimension exist for many\n",
    "  distributions (all in the Gibbs family), the familiar **locality** of the\n",
    "  first-order autoregressive (AR(1)) process—that the conditionals are Markov,\n",
    "  have linear expectation (w.r.t. the neighboring value), and constant variance (idem)—\n",
    "  requires Gaussianity in higher dimensions.\n",
    "  To be clear, the 1D field defined by a simple (and physically interpretable) AR(1) process\n",
    "  $x_{k+1} = a x_k + \\epsilon_k$ exhibits these properties for *any* distribution by construction.\n",
    "  But in 2D or higher-dimensional lattices, such \"directed autoregressive\" constructions become\n",
    "  anisotropic and dependent on the chosen ordering, losing stationarity.\n",
    "  Instead, the AR(1) properties can be salvaged by imposing them directly on the conditionals but,\n",
    "  as shown by Besag (1974), Gaussianity is required to consistently define a global (joint) distribution.\n",
    "  The Markov (sparse) structure also makes the sampling and conditioning even more efficient.\n",
    "  - - -\n",
    "</details>\n",
    "and so will be assumed in what follows.\n",
    "Recall that Gaussian random variables (including vectors and fields)\n",
    "are fully specified by their mean and covariance.\n",
    "Thus, once we are in possession of a covariance matrix,\n",
    "we can sample discrete fields as random vectors, just as we did previously."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "49171bc8",
   "metadata": {},
   "outputs": [],
   "source": [
    "sample_GM, = import_from_nb(\"T2\", [\"sample_GM\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2650758c",
   "metadata": {},
   "source": [
    "Let's try it out.\n",
    "\n",
    "### 1D example"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "58b8bd6d",
   "metadata": {},
   "outputs": [],
   "source": [
    "@interact(**vg_params, N=(1, 30), bottom=list(vg_params))\n",
    "def sample_1D(power=1.5, transf01=\"expo\", Range=0.3, nugget=1e-2, sill=1, N=10):\n",
    "    variogram = variograms(power, transf01, Range, nugget, sill)\n",
    "    C = covar_modelled(grid1D[:, None], variogram)\n",
    "    fields = sample_GM(C=C, N=N, reg=1e-12, rng=3000)\n",
    "\n",
    "    fig, axs = plt.subplots(figsize=(9, 4), ncols=2)\n",
    "    colorbar(fig, axs[0].matshow(C, cmap=\"inferno\", vmin=0, vmax=sill), axs[0])\n",
    "    axs[0].grid(False)\n",
    "    axs[1].plot(grid1D, fields)\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "40184795",
   "metadata": {},
   "source": [
    "#### Exc – impact of variogram parameters\n",
    "\n",
    "Refer to both of the above interactive widgets to answer the following.\n",
    "\n",
    "- (a) What happens to the covariance matrix and the fields when `Range` $\\rightarrow \\infty$ ?\n",
    "- (b) What about `Range` $\\rightarrow 0$ ?\n",
    "- (c) Can you reproduce this using `nugget` somehow ?\n",
    "- (d) Try `nugget = 0` and `0.01`. Explain why the impact is much more noticeable\n",
    "  in the Gauss and Cauchy cases.\n",
    "- (e) How does changing `sill` affect the fields? Can you express `sill` in terms of $C(d)$ somehow? Thus, what other name could we give it?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7a1db7be",
   "metadata": {},
   "outputs": [],
   "source": [
    "# show_answer('variogram params')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3dbc0c69",
   "metadata": {},
   "source": [
    "### 2D example"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e4e15218",
   "metadata": {},
   "outputs": [],
   "source": [
    "vg = variograms(2, \"expo\", Range=1, nugget=0, sill=1)\n",
    "C = covar_modelled(grid2D, vg)\n",
    "C.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "def62f64",
   "metadata": {},
   "source": [
    "Note that the size of the covariance matrix is the square of a single 2D field's size.\n",
    "It's interesting to inspect the covariance matrix in this 2D case,\n",
    "which has a nested structure."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7d6457e9",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = freshfig(\"2D covar\")\n",
    "C0 = vg(np.inf)\n",
    "ax.matshow(C, cmap=\"inferno\", vmin=0, vmax=C0);\n",
    "ax.grid(False)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cb7aafea",
   "metadata": {},
   "source": [
    "Now let's generate and plot some realizations of the corresponding Gaussian\n",
    "random field."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "444c6305",
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot2d(ax, field, contour=True, show_obs=True, cmap=\"PiYG\"):\n",
    "    vmin = -3.5*np.sqrt(C0)\n",
    "    vmax = +3.5*np.sqrt(C0)\n",
    "\n",
    "    ax.set(aspect=\"equal\", xticks=[0, 1], yticks=[0, 1])\n",
    "\n",
    "    if show_obs:\n",
    "        ax.plot(*obs_locations.T, \"ko\", ms=4)\n",
    "        ax.plot(*obs_locations.T, \"yo\", ms=1)\n",
    "\n",
    "    field = field.reshape(grid2x.shape)\n",
    "    if contour:\n",
    "        ε = 1e-12 # fudge to center colors right below 0\n",
    "        levels = np.arange(vmin - ε, vmax + .5 + ε, .5)\n",
    "        return ax.contourf(field, levels=levels, extent=(0, 1, 0, 1), extend=\"both\", cmap=cmap)\n",
    "    else:\n",
    "        return ax.pcolormesh(grid2x, grid2y, field, shading=\"auto\", vmin=vmin, vmax=vmax, cmap=cmap)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "128e0d7e",
   "metadata": {},
   "source": [
    "Beware that `contourf` (which works with regularly spaced grids) uses simple bilinear interpolation\n",
    "between nodes, making the fields appear higher resolution than they actually are.\n",
    "Therefore, you may want to use `contour=False` instead to view the actual grid \"pixels\"."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4df29848",
   "metadata": {},
   "outputs": [],
   "source": [
    "@interact(**vg_params)\n",
    "def sample_2D(power=1.5, transf01=\"expo\", Range=0.3, nugget=1e-2):\n",
    "    fig, axs = plt.subplots(figsize=(8, 4), nrows=2, ncols=4, sharex=True, sharey=True)\n",
    "    C = covar_modelled(grid2D, variograms(power, transf01, Range, nugget, sill=1))\n",
    "    fields = sample_GM(C=C, N=len(axs.ravel()), reg=1e-12, rng=3000)\n",
    "\n",
    "    for ax, field in zip(axs.ravel(), fields.T):\n",
    "        cb = plot2d(ax, field, show_obs=False)\n",
    "\n",
    "    fig.tight_layout()\n",
    "    colorbar(fig, cb, axs)\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2a35b65b",
   "metadata": {},
   "source": [
    "## Estimation\n",
    "\n",
    "Denote $x = X(s)$ the random field at a single location,\n",
    "$\\vect{x}$ the random field at all grid locations,\n",
    "and $\\vect{y}$ the field at the observation locations.\n",
    "\n",
    "### Problem (2D)\n",
    "\n",
    "For our estimation/interpolation target ($\\vect{x}$) we sample using the default covariance matrix defined above."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d7c7da6a",
   "metadata": {},
   "outputs": [],
   "source": [
    "truth = sample_GM(C=C, N=1, reg=1e-12).squeeze()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "42ae9551",
   "metadata": {},
   "source": [
    "For our observations ($\\vect{y}$) we randomly pick some grid points:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a15a0ef6",
   "metadata": {},
   "outputs": [],
   "source": [
    "nObs = 10\n",
    "rnd.seed(3000)\n",
    "obs_indices = rnd.choice(len(grid2D), nObs, replace=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b82e76e9",
   "metadata": {},
   "source": [
    "...such that the observation locations and values are simply:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "714bf03f",
   "metadata": {},
   "outputs": [],
   "source": [
    "obs_locations = grid2D[obs_indices]\n",
    "observations = truth[obs_indices]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a08e120b",
   "metadata": {},
   "source": [
    "However (unlike classical bilinear/bicubic interpolation),\n",
    "the following methods work perfectly well with irregular observations,\n",
    "irregular grids, or different orderings of the flattened grid nodes.\n",
    "\n",
    "### Methods\n",
    "\n",
    "The following is a register of fields to include in plot below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4ddc6a71",
   "metadata": {},
   "outputs": [],
   "source": [
    "estims = dict(Truth=truth)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a0891244",
   "metadata": {},
   "source": [
    "Each time you (have tried to) implement one of the following methods,\n",
    "come back up here and re-run this next cell\n",
    "to visually compare and assess your method."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "deca3b12",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axs = freshfig(figsize=(8, 4), squeeze=False, nrows=2, ncols=len(estims), sharex=True, sharey=True)\n",
    "axs[0, 0].set_ylabel(\"Estimate\");\n",
    "axs[1, 0].set_ylabel(\"Errors\");\n",
    "\n",
    "for name, ax1, ax2 in zip(estims, *axs):\n",
    "    x = estims[name]\n",
    "    ax1.set(title=name)\n",
    "    cb1 = plot2d(ax1, x)\n",
    "    cb2 = plot2d(ax2, x - truth, cmap=\"seismic\")\n",
    "\n",
    "fig.tight_layout()\n",
    "colorbar(fig, cb1, axs[0]);\n",
    "colorbar(fig, cb2, axs[1]);"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "73f673b5",
   "metadata": {},
   "source": [
    "Pre-compute some objects that see repeated use."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f5370062",
   "metadata": {},
   "outputs": [],
   "source": [
    "dists_yy = dist_euclid(obs_locations, obs_locations)\n",
    "dists_xy = dist_euclid(grid2D, obs_locations)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "26c00359",
   "metadata": {},
   "source": [
    "The cells below contain snippets of different spatial interpolation methods\n",
    "followed by a cell that plots the interpolants.\n",
    "Complete the code snippets.\n",
    "\n",
    "#### Exc: Nearest neighbour interpolation\n",
    "\n",
    "Implement the method [(Wikipedia)](https://en.wikipedia.org/wiki/Nearest-neighbor_interpolation)\n",
    "by adjusting the code below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a5e8cf27",
   "metadata": {},
   "outputs": [],
   "source": [
    "nearest_obs = np.zeros(len(grid2D), dtype=int)  ### FIX THIS ###\n",
    "estims[\"Nearest-n.\"] = observations[nearest_obs]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4c9c012f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# show_answer('nearest neighbour interp')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b1f2d7df",
   "metadata": {},
   "source": [
    "#### Exc: Inverse distance weighting\n",
    "\n",
    "Implement the method [(Wikipedia)](https://en.wikipedia.org/wiki/Inverse_distance_weighting).\\\n",
    "*Hint: The `errstate` context silences warnings due to division by 0 (whose special case is treated further down).*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e21fc64b",
   "metadata": {},
   "outputs": [],
   "source": [
    "exponent = 3\n",
    "with np.errstate(invalid='ignore', divide='ignore'):\n",
    "    weights = np.zeros_like(dists_xy)  ### FIX THIS ###    \n",
    "estims[\"Inv-dist.\"] = weights @ observations # Apply weights\n",
    "estims[\"Inv-dist.\"][obs_indices] = observations # Fix singularities"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9d22ec65",
   "metadata": {},
   "outputs": [],
   "source": [
    "# show_answer('inv-dist weight interp')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "84e46ea4",
   "metadata": {},
   "source": [
    "#### Kriging\n",
    "\n",
    "Kriging finds the best (minimum) mean square error (MSE), $\\Expect (\\widehat{x} - x)^2$, among all linear \"predictors\",\n",
    "$\\widehat{x} = \\vect{w}\\tr \\vect{y}$, that are unbiased (BLUP).\n",
    "\n",
    "<a name='Exc-–-\"simple\"-kriging-(SK)'></a>\n",
    "\n",
    "#### Exc – \"simple\" kriging (SK)\n",
    "\n",
    "Suppose $X(s)$ has a mean that is constant in space, $\\mu$, and known.\n",
    "Since it is easy to subtract (and later re-include) $\\mu$\n",
    "from both $x$ and the data $\\vect{y}$, we simply assume $\\mu = 0$.\n",
    "Thus, $\\Expect \\vect{y} = \\vect{0}$ and $\\Expect \\widehat{x} = 0$ for any weights $\\vect{w}$,\n",
    "and so $\\widehat{x}$ is already unbiased.\n",
    "Meanwhile,\n",
    "$$\n",
    "\\begin{align}\n",
    "  \\text{MSE}\n",
    "  % \\Expect \\big( \\widehat{x} - x \\big)^2\n",
    "  &= \\Expect \\big( \\vect{w}\\tr \\vect{y} - x \\big)^2 \\\\\n",
    "  % &= \\Expect \\big( \\vect{w}\\tr \\vect{y} \\vect{y}\\tr \\vect{w} - 2 x \\vect{y}\\tr \\vect{w} + x^2 \\big) \\\\\n",
    "  &= \\vect{w}\\tr \\Expect \\big( \\vect{y} \\vect{y}\\tr \\big) \\vect{w}\n",
    "     - 2 \\Expect \\big( x \\vect{y}\\tr \\big) \\vect{w}\n",
    "     + \\Expect \\big( x^2 \\big) \\\\\n",
    "  &= \\vect{w}\\tr \\mat{C}_{\\vect{y} \\vect{y}} \\vect{w}\n",
    "     - 2 \\vect{c}_{x \\vect{y}}\\tr \\vect{w}\n",
    "     + C(0) \\,,\n",
    "\\end{align}\n",
    "$$\n",
    "where $C(0)$ is the variance (\"sill\") of the field.\n",
    "The minimum can be found by setting the derivative with respect to $\\vect{w}$ to zero,\n",
    "yielding $\\vect{w}_{\\text{SK}} = \\mat{C}_{\\vect{y} \\vect{y}}^{-1} \\, \\vect{c}_{x \\vect{y}}$\n",
    "\n",
    "<details style=\"border: 1px solid #aaaaaa; border-radius: 4px; padding: 0.5em 0.5em 0;\">\n",
    "  <summary style=\"font-weight: normal; font-style: italic; margin: -0.5em -0.5em 0; padding: 0.5em;\">\n",
    "    We have already seen this formula in the\n",
    "    <a href=\"T5%20-%20Multivariate%20Kalman%20filter.ipynb#Regression-perspective\">regression perspective on the Kalman gain (T5)</a>,\n",
    "    which also noted the Gauss-Markov theorem which states that it is optimal in the BLUE sense.\n",
    "    The perspective of radial basis functions (RBF) is also informative (optional reading...)\n",
    "  </summary>\n",
    "\n",
    "  The BLUP criterion is similar to the best linear unbiased *estimator* (BLUE),\n",
    "  and so the equivalence of the formulae may not come as a surprise.\n",
    "  Still, even though the MSE optimized by the BLUP can be construed as a *variance* (by the unbiasedness constraint),\n",
    "  it should not be confused with the variance minimized by the BLUE, which targets a *fixed* (but unknown) quantity\n",
    "  (one conventionally \"estimates\" fixed effects and \"predicts\" random effects).\n",
    "  Indeed, fruitfully applying the BLUE criterion to the kriging setting is difficult.\n",
    "  For example, the additional randomness (i.e., in the estimand $x$ as well as in the observations $\\vect{y}$)\n",
    "  makes it so that if the mean $\\mu$ is known – effectively zero – then the BLUE is simply zero\n",
    "  (it is unbiased and has zero variance)!\n",
    "  Or, if $\\mu$ is unknown, we can assume $\\vect{y} = x + \\vect{r}$ in order to define unbiasedness,\n",
    "  which then requires that the weights sum to one, i.e. $\\vect{w}\\tr\\vect{1} = 1$.\n",
    "  Then, if $\\mat{C}_{\\vect{y}\\vect{y}}$ is the covariance of $\\vect{r}$,\n",
    "  the BLUE weights become $(\\vect{1}\\tr \\mat{C}_{\\vect{y}\\vect{y}}^{-1})/(\\vect{1}\\tr \\mat{C}_{\\vect{y}\\vect{y}}^{-1} \\vect{1})$,\n",
    "  and so $\\widehat{x}(s)$ is constant in $s$, i.e. flat.\n",
    "  Thus we see the need to include the randomness of $x$ along with that of $\\vect{y}$.\n",
    "  This seems contrary to the classical BLUE framing,\n",
    "  but we have already seen it done for the Kalman gain (by augmenting the observations by the prior mean).\n",
    "  Even so, claims of BLUE optimality would be limited to cases where we could find\n",
    "  covariances and operators that fit with the variogram-modelled covariance,\n",
    "  which is very limiting.\n",
    "\n",
    "  Now for the \"dual\" perspective of kriging, which is held by **radial basis function (RBF) interpolation**.\n",
    "  Ultimately, kriging (simple, ordinary, and universal) provides an estimate of the form\n",
    "  $\\widehat{x} = \\vect{y}\\tr (\\mat{A}_\\vect{y}^{-1} \\vect{b}_{\\vect{y} x})$,\n",
    "  where $\\vect{y}$ is the observations,\n",
    "  and the subscripts indicate that $\\mat{A}_\\vect{y}$ depends on the locations of $\\vect{y}$,\n",
    "  while $\\vect{b}_{\\vect{y} x}$ depends on the locations of $\\vect{y}$ and $x$ both.\n",
    "  In the kriging perspective, $\\vect{w} = \\mat{A}_\\vect{y}^{-1} \\vect{b}_{\\vect{y} x}$ is seen as the \"weights\" and get applied to $\\vect{y}$:\n",
    "  once solved for, they are easily re-applied for different values of $\\vect{y}$.\n",
    "  However, if we shift the parentheses (and computational order) onto the first two factors,\n",
    "  then the linear solver computes $\\vect{v} = \\mat{A}_\\vect{y}^{-1} \\vect{y}$,\n",
    "  fitting $N$ RBFs to $N$ data points, $\\vect{y}$.\n",
    "  Now the \"weights\", $\\vect{v}$, are computed from the observations,\n",
    "  but applied onto $\\vect{b}_{\\vect{y} x}$, which is a function of (the location) of $x$.\n",
    "  As such we have not merely derived an estimate for a single location, $s$,\n",
    "  but an *interpolant* (that is cheap to evaluate) for all locations.\n",
    "\n",
    "  Lastly, why is kriging said to be **non-parametric**?\n",
    "  Heuristically, note that – while the kernel itself is specified through a small number of parameters –\n",
    "  kriging defines a separate linear regression problem for each and every location of the domain,\n",
    "  of which there are infinitely (uncountably) many.\n",
    "  More rigorously, by identifying the variogram with a kernel,\n",
    "  we can appeal to the theory of reproducing kernel Hilbert spaces (RKHS) to kriging.\n",
    "  In particular, kriging (GP) can be seen to solve an \"empirical risk minimisation problem\"\n",
    "  with negligible measurement error (for interpolation).\n",
    "  Thus, by the \"representer theorem\",\n",
    "  – even though we're only solving for $N$ weights, $\\vect{v}$ –\n",
    "  the fitted surface is not restricted to a finite feature basis,\n",
    "  but is the minimum in an infinite-dimensional function space.\n",
    "\n",
    "  - - -\n",
    "</details>\n",
    "\n",
    "However, we previously derived this $\\widehat{x}$ as the posterior/conditional mean of Gaussian distribution.\n",
    "This perspective is that of GP regression,\n",
    "but was also highlighted by [Krige (1951)](#References),\n",
    "and makes it evident that kriging also provides an uncertainty estimate\n",
    "(a major benefit compared to the other interpolation methods we tested above),\n",
    "namely the posterior covariance matrix, $C(0) - \\vect{c}_{x \\vect{y}}\\tr \\mat{C}_{\\vect{y} \\vect{y}}^{-1} \\vect{c}_{x \\vect{y}}$.\n",
    "\n",
    "Implement the method [(Wikipedia)](https://en.wikipedia.org/wiki/Kriging#Simple_kriging).  \n",
    "*Hint: You may use `sla.inv`, but `sla.solve` is better, and `sla.lstsq` is even better.*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0ab70a46",
   "metadata": {},
   "outputs": [],
   "source": [
    "C0 = vg(np.inf)\n",
    "### FIX THIS ###\n",
    "covar_yy = ...\n",
    "cross_xy = ...\n",
    "weights = np.zeros_like(dists_xy.T) # cross_yy \"\\\" covar_yy\n",
    "estims[\"S.Kriging\"] = observations @ weights"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d5648560",
   "metadata": {},
   "outputs": [],
   "source": [
    "# show_answer('Simple kriging')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "80e19eb6",
   "metadata": {},
   "source": [
    "### More kriging (1D)\n",
    "\n",
    "The above 2D case used the same variogram that generated the truth.\n",
    "But in practice, we do not know this variogram (whose very existence is a theoretical assumption),\n",
    "so it must also be estimated.\n",
    "To that end, let's gain some understanding of the variogram's impact\n",
    "on the resulting *estimate*. We focus on the 1D problem because its illustrations are clearer.\n",
    "\n",
    "For our estimation target, we use $x(s) = \\sin (s^2)$ for $s \\in [0, L]$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b838f4a6",
   "metadata": {},
   "outputs": [],
   "source": [
    "def true(s):\n",
    "    return np.sin(s**2)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c2d2013f",
   "metadata": {},
   "source": [
    "The observations are taken at the following (configurable) locations:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a2c07f60",
   "metadata": {},
   "outputs": [],
   "source": [
    "def gen_obs_loc(nObs, spacing, L):\n",
    "    return L/2 + spacing * np.linspace(-L/2, L/2, nObs)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8b6c221d",
   "metadata": {},
   "source": [
    "Visualisation:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2ea61cdb",
   "metadata": {},
   "outputs": [],
   "source": [
    "@interact(**vg_params, L=(1, 10, 0.1), nObs=(1, 100, 1), spacing=(0.01, 1), bottom=list(vg_params))\n",
    "def kriging_1d(L=4, nObs=6, spacing=0.5, transf01=\"expo\", power=1.5, Range=0.3, nugget=0, sill=1):\n",
    "    # Experiment setup\n",
    "    sx = np.linspace(0, L, 1001)       # locations of x (interpolation grid)\n",
    "    sy = gen_obs_loc(nObs, spacing, L) # locations of y (observations points)\n",
    "    truth = true(sx)\n",
    "    obsvs = true(sy)\n",
    "\n",
    "    # Plot setup\n",
    "    fig, ax = plt.subplots(figsize=(8, 4))\n",
    "    ax.set_ylim(-2, 2)\n",
    "    ax.plot(sx, truth, \"-k\", label=\"truth\")\n",
    "    ax.plot(sy, obsvs, 'ok', label='data', ms=12, zorder=9)\n",
    "\n",
    "    # Kriging setup\n",
    "    vg = variograms(power, transf01, Range=Range, nugget=nugget, sill=sill)\n",
    "    dists_yy = dist_euclid(sy, sy)\n",
    "    dists_xy = dist_euclid(sx, sy)\n",
    "\n",
    "    with nonchalance():\n",
    "        mu = 0\n",
    "        ax.axhline(mu, c=\"k\", lw=0.5)\n",
    "        interp = mu + simple_kriging(vg, dists_xy, dists_yy, obsvs - mu)\n",
    "        ax.plot(sx, interp, 'C1', label=\"S-krig\", lw=5)\n",
    "\n",
    "    ax.legend(loc='lower left', ncols=2)\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "be161864",
   "metadata": {},
   "source": [
    "#### Exc: Impact of variogram on kriged fields\n",
    "\n",
    "Add the simple kriging interpolant by copy-pasting your solution from above into the function below.\n",
    "*Hint: Move (and de-indent) code out of the `nonchalance()` context if you want to be\n",
    "able to view error messages (to debug errors).*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "26ce9578",
   "metadata": {},
   "outputs": [],
   "source": [
    "def simple_kriging(vg, dists_xy, dists_yy, observations, mu):\n",
    "    field_estimate = ...\n",
    "    return field_estimate"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "89ab959b",
   "metadata": {},
   "source": [
    "- (a) What happens when `Range` $\\to 0$ ? What about $\\to \\infty$? What about intermediate values?\n",
    "- (b) Try the same with Gauss/Cauchy variograms. What is the main difference?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "48d944b1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# show_answer('Interpolant = f(Variogram)', 'a')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dcd6ccdd",
   "metadata": {},
   "source": [
    "#### Exc (optional) – Ordinary kriging\n",
    "\n",
    "Let us do away with the assumption that the mean of the field is known,\n",
    "all the while retaining the assumption that it is constant in space.\n",
    "The resulting method is called ordinary kriging.\n",
    "In this case, unbiasedness of $\\widehat{x} = \\vect{w}\\tr \\vect{y}$ requires that the weights sum to one.\n",
    "This can be imposed on the MSE minimization using a Lagrange multiplier $\\lambda$,\n",
    "yielding the augmented system to solve:\n",
    "$$ \\begin{pmatrix} \\mat{C}_{\\vect{y} \\vect{y}} & \\vect{1} \\\\ \\vect{1}\\tr & 0 \\end{pmatrix} \\begin{pmatrix} \\vect{w} \\\\ \\lambda \\end{pmatrix}\n",
    "= \\begin{pmatrix} \\vect{c}_{x \\vect{y}} \\\\ 1 \\end{pmatrix} \\,.$$\n",
    "\n",
    "Ordinary kriging can be reproduced by separately kriging the mean and the resulting residual field [(Wackernagel, 2013)](#References).\n",
    "\n",
    "It is also important to note that the MSE – by its resemblance to intrinsic increments –\n",
    "can be expressed in terms of variograms even when the covariance function does not exist,\n",
    "and the same goes for the method of OK. Somewhat surprisingly, the linear system to be solved looks exactly the same, except with covariances replaced by variograms.\n",
    "In this case, the unbiasedness requirement also manifests as a requirement\n",
    "that any variance expressed as a linear combination of variograms be positive.\n",
    "In the RBF perspective, this requirement guarantees the existence and uniqueness of the solution to the OK equations.\n",
    "\n",
    "- (a) Implement the OK method and add it to the above plotting by adding the following to its code:\n",
    "  ```python\n",
    "  with nonchalance():\n",
    "    interp = ordinary_kriging(vg, dists_xy, dists_yy, obsvs)\n",
    "    ax.plot(sx, interp, 'C2', label=\"O-krig\", lw=4)\n",
    "  ```\n",
    "- (b) Which variogram model(s) does OK \"unlock\"?\n",
    "- (c) What is peculiar about the interpolant for each of the new variograms?\n",
    "- (d) Set `nObs = 1`. What is the difference between the SK and OK interpolants?\n",
    "- (e) Set `nObs = 2`. What is the difference between the SK and OK interpolants? Explore with a few different `spacing` values."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b975a812",
   "metadata": {},
   "outputs": [],
   "source": [
    "# show_answer('Ordinary kriging', 'a')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "03015c86",
   "metadata": {},
   "source": [
    "#### Exc (optional) – Universal kriging\n",
    "\n",
    "In addition to the flat/constant *feature* $\\vect{1}$,\n",
    "one can add additional regressors in a similar fashion to how the OK system augments the SK system of equations,\n",
    "producing universal kriging (UK).\n",
    "If the features are monomials, then they may well be called *trends*,\n",
    "and the UK method is called intrinsic kriging (IK).\n",
    "Their order $+1$ defines the maximum allowable order of the *generalized* covariance function\n",
    "(for example, `power=3` \"variogram\" has order 2, requiring order 1 (linear) trends),\n",
    "again rendering the augmented linear system well-posed."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "247dee32",
   "metadata": {},
   "outputs": [],
   "source": [
    "order = 1\n",
    "features = [(lambda x: x**p) for p in range(order+1)]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "88bd550e",
   "metadata": {},
   "source": [
    "- (a) Add the following to the plotting code above, and implement the requisite UK method.\n",
    "  ```python\n",
    "    with nonchalance():\n",
    "        regressors = [f(sy) for f in features]\n",
    "        regressands = [f(sx) for f in features]\n",
    "        interp = universal_kriging(vg, dists_xy, dists_yy, obsvs, regressors, regressands)\n",
    "        ax.plot(sx, interp, 'C3', label=\"U-krig\", lw=3)\n",
    "  ```\n",
    "- (b) Add the interpolation method of cubic spine (from scipy; Google it!).\\\n",
    "  Can you make the UK interpolant reproduce this (in the interior)?\n",
    "  What happens when you vary `Range`?\n",
    "- (c) Although theoretically suboptimal\n",
    "  (since it assumes a constant/flat mean, and does not support generalized covariance functions),\n",
    "  the OK system is still solvable with a cubic variogram.\n",
    "  Can you find experimental control settings that show the OK estimate does not\n",
    "  exactly reproduce the spline solution?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2e052bbe",
   "metadata": {},
   "outputs": [],
   "source": [
    "# show_answer('Universal kriging', 'a')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "04912618",
   "metadata": {},
   "source": [
    "Further generalizations include co-kriging (vector-valued fields).\n",
    "Also, since the kriging mean is smoother than any realization (has lower variance,\n",
    "which favors but does not guarantee spatial/spectral smoothness),\n",
    "conditional generation (simulation) of fields is frequently used.\n",
    "\n",
    "## Summary\n",
    "\n",
    "Random, (geo)spatial fields are often modeled by their (auto-)correlation function or\n",
    "– more generally – the (semi-)variogram.\n",
    "Covariances also form the basis of a family of spatial interpolation and approximation\n",
    "methods known as kriging.\n",
    "\n",
    "### Next: [T8 - Monte-Carlo & ensembles](T8%20-%20Monte-Carlo%20%26%20cov%20estimation.ipynb)\n",
    "\n",
    "<a name=\"References\"></a>\n",
    "\n",
    "### References\n",
    "\n",
    "<!--\n",
    "\n",
    "@book{chiles2012geostatistics,\n",
    "  title={Geostatistics: Modeling Spatial Uncertainty},\n",
    "  author={Chil{\\`e}s, J.P. and Delfiner, P.},\n",
    "  isbn={9780470317839},\n",
    "  series={Wiley Series in Probability and Statistics},\n",
    "  url={https://books.google.no/books?id=tZl07WdjYHgC},\n",
    "  year={2012},\n",
    "  publisher={Wiley}\n",
    "}\n",
    "\n",
    "@book{wackernagel2013multivariate,\n",
    "  title={Multivariate Geostatistics: An Introduction with Applications},\n",
    "  author={Wackernagel, H.},\n",
    "  isbn={9783662052945},\n",
    "  year={2013},\n",
    "  publisher={Springer Berlin Heidelberg}\n",
    "}\n",
    "\n",
    "@article{krige1951statistical,\n",
    "  title={A statistical approach to some basic mine valuation problems on the Witwatersrand},\n",
    "  author={Krige, Daniel G},\n",
    "  journal={Journal of the Southern African Institute of Mining and Metallurgy},\n",
    "  volume={52},\n",
    "  number={6},\n",
    "  pages={119--139},\n",
    "  year={1951},\n",
    "  publisher={Southern African Institute of Mining and Metallurgy}\n",
    "}\n",
    "-->\n",
    "\n",
    "- **Chilès & Delfiner (2012)**:\n",
    "  J.P. Chilès and P. Delfiner,\n",
    "  \"Geostatistics: Modeling Spatial Uncertainty\", 2nd edition,\n",
    "  Wiley, 2012.\n",
    "- **Wackernagel (2013)**:\n",
    "  H. Wackernagel, \"Multivariate Geostatistics: An Introduction with Applications\", Springer, 2013.\n",
    "- **Krige (1951)**:\n",
    "  D. G. Krige, \"A statistical approach to some basic mine valuation problems on the Witwatersrand\",\n",
    "  1951."
   ]
  }
 ],
 "metadata": {
  "jupytext": {
   "cell_metadata_filter": "-all",
   "formats": "ipynb,nb_mirrors//py:light,nb_mirrors//md"
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
